/*!
    Copright: Kai "DrHalan" Mast - www.freakybytes.org

    This file is part of OpenNetwork. See OpenNetwork.h for license information.
*/

#include "CTCPPort.h"

#include "include/OpenNetwork/OpenNetwork.h"

#include <errno.h>

namespace onet
{

    CTCPPort::CTCPPort(INetworkCallback* callback_, u16 maxConnections_, u16 port_, bool listen_, bool IPv6_)
            : INetworkPort(callback_, maxConnections_, port_, listen_, IPv6_), peerCount(0), highest_fd(0), listenfd(-1)
    {
        if (maxPeers)
            peers.init(maxPeers);

        FD_ZERO(&peer_fds);
    }

    CTCPPort::~CTCPPort()
    {
        this->disconnectAll();

        //! Important to close the file descriptor afterwards
        if (this->isListening())
        {
            if (close(listenfd) == -1)
                perror("close");
        }
    }

    bool CTCPPort::init()
    {
        //! If this port isn't listening for connections we don't even need a listen-filedscriptor
        if (this->isListening())
        {
            //! The actual network code
            //! We've to check wheter the port we use uses IPv6 or not and differentiate
            if (this->isIPv6())
                listenfd = socket(PF_INET6, SOCK_STREAM, 0);
            else
                listenfd = socket(PF_INET, SOCK_STREAM, 0);

            if (listenfd == -1)
                return false;

            const s32 y = 1;

#ifndef WINDOWS
            setsockopt( listenfd, SOL_SOCKET, SO_REUSEADDR, &y, sizeof(s32) );
#else
            setsockopt( listenfd, SOL_SOCKET, SO_REUSEADDR, (char*)&y, sizeof(s32) );
#endif

            if (this->isIPv6())
            {
                struct sockaddr_in6 myaddr;

                myaddr.sin6_family = AF_INET6;

                myaddr.sin6_port = htons( portNumber );
                myaddr.sin6_addr = in6addr_any;

                if (bind(listenfd, (struct sockaddr*)&myaddr, sizeof(myaddr)) == -1)
                {
                    perror("bind");
                    return false;
                }
            }
            else
            {
                struct sockaddr_in myaddr;

                myaddr.sin_family = AF_INET;

                myaddr.sin_port = htons( portNumber );
                myaddr.sin_addr.s_addr = INADDR_ANY;

                if (bind(listenfd, (struct sockaddr*)&myaddr, sizeof(myaddr)) == -1)
                {
                    perror("bind");
                    return false;
                }
            }

            //! Make the port listening for connections
            if (::listen( listenfd, 20) != 0)
            {
                perror("listen");
                return false;
            }
        }

        this->getCallback()->onCreate(this);

        return true;
    }

    SNetworkAddress CTCPPort::getAddress(const PeerID id)
    {
        if (peers[id])
        {
            if (this->isIPv6())
                return SNetworkAddress("::/128", 0, true);
            else
                return SNetworkAddress("0.0.0.0", 0, false);
        }

        return peers[id]->address;
    }

//! Connect to an address
//! Returns false if already connected or some error occured
    bool CTCPPort::connect(const SNetworkAddress& address)
    {
        if (this->isConnected(address))
            return false;

        if (address.IPv6 != this->isIPv6())
            return false;

        for (u16 u = 0; u < this->getMaxPeers(); u++)
        {
            if (!peers[u])
            {
                //! Create the socket
                s32 socketFd = -1;

                if (this->isIPv6())
                    socketFd = socket(PF_INET6, SOCK_STREAM, 0);
                else
                    socketFd = socket(PF_INET, SOCK_STREAM, 0);

                if (socketFd < 0)
                {
                    perror("socket");
                    return false;
                }

                peers.setValue(u, new STCPPeer(address, socketFd));

                this->addPeer(u);

                return true;
            }
        }

        return false;
    }

//! The disconnect pendant
//! Returns false if the peerID is invalid
    bool CTCPPort::disconnect(const PeerID id)
    {
        if (id < 0)
            return false;

        if (!peers[id])
            return false;

        if (close(peers[id]->socketfd) == -1)
            perror("close");

        this->removePeer(id);

        return true;
    }

//! Disconnects all peers
//! Returns number of peers disconnected
    u16 CTCPPort::disconnectAll(void)
    {
        u32 count = 0;

        for (u16 u = 0; u < maxPeers; u++)
        {
            count += this->disconnect(u);
        }

        return count;
    }

    bool CTCPPort::isConnected(const SNetworkAddress& addr)
    {
        for (u16 u = 0; u < maxPeers; u++)
        {
            if (peers[u] && peers[u]->address == addr)
                return true;
        }

        return false;
    }

//! Gets the Id of a certain address
    PeerID CTCPPort::getPeer(const SNetworkAddress& addr)
    {
        for (u16 u = 0; u < maxPeers; u++)
        {
            if (peers[u] && peers[u]->address == addr )
                return u;
        }

        return -1;
    }

//! Sends and receives data
    void CTCPPort::work(void)
    {
        this->getCallback()->onUpdate();

        //! Clear the sets
        FD_ZERO(&read_fds);
        FD_ZERO(&write_fds);

        //! First we add all peers to the reading set and the writing set
        read_fds = write_fds = peer_fds;

        //! We set the timeout to 10 milliseconds
        timeval time;
        time.tv_sec = 0;
        time.tv_usec = 10000;

        //! Now check what filedescriptors are free
        s32 numfd = highest_fd;

        //! If were listening
        if (this->isListening())
        {
            //! also the listening port
            FD_SET(listenfd, &read_fds);

            if (listenfd > highest_fd)
                numfd = listenfd;
        }

        if (select(numfd+1, &read_fds, &write_fds, 0, &time) == -1)
            perror("select");
        else
        {
            //! Check for new connections
            //! If we're listening
            if (this->isListening() && FD_ISSET(listenfd, &read_fds) && (peerCount <= maxPeers))
                    this->processIncomingConnection();

            //! Check for new messages
            for (u32 u = 0; u < maxPeers; u++)
            {
                if (peers[u])
                {
                    if (peers[u]->connected)
                    {
                        if (FD_ISSET(peers[u]->socketfd,&read_fds))
                            this->processReceiving(u);
                    }
                    else
                    {
                        //! Handle connecting to the peer
                        //! Again we've to check wheter we can write to the socket

                        if (FD_ISSET(peers[u]->socketfd, &write_fds))
                            this->processConnecting(u);
                    }
                }
            }

            this->processSending();
        }
    }

    //! Get the protocol used
    EPortType CTCPPort::getType() const
    {
        return EPT_TransmissionControlProtocol;
    }

    //! Send data to peer - Reliable or not
    bool CTCPPort::send(const PeerID id, const u8 type, const u32 length, const c8* data, const bool reliable)
    {
        //! Message Protocol
        //! 1 Byte - Type (user defined)
        //! 4 Bytes - Length
        //! Length Bytes - Data

        //! -> length + 5 Bytes

        //! TCP only supports reliable messageing
        //! We decided to ignore that flag right away so you can easily switch between UDP/TCP
        //! ----------
        //! if(!reliable)
        //!    return ESMR_NotSupported;

        if (!peers[id])
            return false;

        if (!peers[id]->connected)
            return false;

        BitStream *message = new BitStream(length + 5);

        //! we add the type given by the user and the data
        message->write<u8>(type, false);
        message->write<u32>(length, false);

        //! You can also send empty messages (e.g. ping/pong)
        if (length > 0)
            message->writeRawArray<c8>(data, length, false);

        //! Reset the bitStream so CTCPPort::work() can read it properly
        message->setByteIndex(0);

        //! Put the message in the queue to handle the sending of it
        messages.push_back(STCPMessage(id, message));

        return true;
    }

    //! Send data to address
    bool CTCPPort::sendRaw(const SNetworkAddress& address, const u16 length, const c8* data)
    {
        return false;
    }

    //! Query certain features
    bool CTCPPort::hasFeature(EFeatureType type) const
    {
        if (type == EFT_RawMessage)
            return false;
        else if (type == EFT_UnreliableMessage)
            return false;
        else if (type == EFT_ReliableMessage)
            return true;
        else if (type == EFT_Latency)
            return false;
        else if (type == EFT_BigMessage)
            return true;
        else
            return false;
    }

    u16 CTCPPort::getPeerCount()
    {
        return peerCount;
    }

    //! Remove a peer (!= disconnect)
    void CTCPPort::removePeer(PeerID id)
    {
        FD_CLR(peers[id]->socketfd, &peer_fds);

        s32 sfd = peers[id]->socketfd;

        peers.remove(id);

        //! Update the highest filedescriptor?
        if (sfd == highest_fd)
        {
            highest_fd = 0;

            for (u32 u = 0; u < maxPeers; u++)
            {
                if (peers[u] && peers[u]->socketfd > highest_fd)
                    highest_fd = peers[u]->socketfd;
            }
        }

        //! Decrease the peer count
        peerCount--;
    }

    //! Adds a peer (!= connect)
    void CTCPPort::addPeer(PeerID id)
    {
        //!Update the peer set
        FD_SET(peers[id]->socketfd, &peer_fds);

        if (peers[id]->socketfd > highest_fd)
            highest_fd = peers[id]->socketfd;

        //!Increase the count of peers
        peerCount++;
    }

    s32 CTCPPort::getLatency(const PeerID id)
    {
        //! Not implented yet
        return -1;
    }

    void CTCPPort::processSending()
    {
        std::vector<std::vector<STCPMessage>::iterator> removals;

        //! Send messages
        for (std::vector<STCPMessage>::iterator it = messages.begin(); it != messages.end(); it++)
        {
            //! Is this message still valid and can we write?
            if (!peers[(*it).peer])
            {
                delete (*it).data;
                removals.push_back(it);
            }
            else if (FD_ISSET(peers[(*it).peer]->socketfd, &write_fds))
            {
                //! Now lets send via TCP
                u32 pos = (*it).data->getByteIndex();
                u32 readLength = (*it).data->getLength() - pos;
                c8* data_string = (*it).data->readRawArray<c8>(readLength);
                s32 sentLength = ::send(peers[(*it).peer]->socketfd, data_string, readLength, 0);

                delete []data_string;

                if (sentLength < 0)
                    perror("send");
                if (u32(sentLength) >= readLength)
                {
                    //! Only delete if the full data was sent
                    delete (*it).data;
                    removals.push_back(it);
                }
                else
                {
                    //! Move to the posting + sentLength
                    (*it).data->setByteIndex(pos + sentLength);
                }

            }
        }

        //! Remove sended messages
        for (std::vector<std::vector<STCPMessage>::iterator>::iterator it = removals.begin(); it != removals.end(); it++)
        {
            messages.erase(*it);
        }
    }

    void CTCPPort::processReceiving(PeerID id)
    {
        //! Clear the buffer (global member)
        memset(&buffer[0], 0, 1024);

        //! The low level TCP code
        s32 readLength  = recv( peers[id]->socketfd, &buffer[0], 1024, 0);

        if (readLength < 0)
        {
            if(errno == ETIMEDOUT)
                this->getCallback()->onPeerDisconnect(id, ENI_ConnectionTimeout);
            else
            {
                this->getCallback()->onPeerDisconnect(id, ENI_UnknownError);
                perror("recv");
            }

            this->removePeer(id);
        }
        else if (readLength == 0)
        {
            this->getCallback()->onPeerDisconnect(id, ENI_ConnectionClosed);
            this->removePeer(id);
        }
        else
        {
            //! Setup the BitStream
            BitStream data(&buffer[0], readLength);

            //! Loop until we finished reading
            while (data.getByteIndex() < data.getLength())
            {
                if (!peers[id]->cache)
                {
                    //! The first byte is the message Type
                    EMessageType type = static_cast<EMessageType>(data.read<u8>());

                    //! The length of the message
                    u32 length = data.read<u32>();

                    //! Check wheter the whole package was read
                    if (length <= ( data.getLength()-data.getByteIndex() ) )
                    {
                        //! Actual data for the Callback! wohoooo
                        c8* message = data.readRawArray<c8>(length);

                        this->getCallback()->onMessageReceived(id, type, length, message, true);
                    }
                    else
                    {
                        //! Push the data into the peers cache
                        peers[id]->msgSize = length;
                        peers[id]->msgType = type;
                        peers[id]->cache = new BitStream(peers[id]->msgSize);

                        //! Append everything to the cache
                        //! We would like to read lenght-5 but we can only read some of it
                        data.copyRawArray<c8>(peers[id]->cache, data.getLength()-data.getByteIndex());
                    }
                }
                else
                {
                    //! Calculate how much we have to read
                    u32 length = peers[id]->msgSize - peers[id]->cache->getByteIndex();

                    //! If we STILL didn't receive everything we just read a piece of it again
                    if (static_cast<u32>(readLength) < length)
                        length = readLength;

                    //! Append data to the cache :)
                    data.copyRawArray<c8>(peers[id]->cache, length, true, false);

                    if (peers[id]->msgSize <= peers[id]->cache->getByteIndex())
                    {
                        //! Everything was received!
                        this->getCallback()->onMessageReceived(id, peers[id]->msgType, peers[id]->msgSize, peers[id]->cache->getData(), true);

                        //! Clean up
                        delete peers[id]->cache;
                        peers[id]->cache = 0;
                    }
                }
            }
        }
    }

    void CTCPPort::processConnecting(PeerID id)
    {
        if (this->isIPv6())
        {
            sockaddr_in6 tmpAddr;

            if (getSockAddress6(tmpAddr, peers[id]->address) &&  ::connect ( peers[id]->socketfd, (sockaddr*)&tmpAddr , sizeof(tmpAddr)) != -1)
                {
                    peers[id]->connected = true;

                    this->getCallback()->onPeerConnect(id);
                }
        }
        else
        {
            sockaddr_in tmpAddr;

            if (getSockAddress(tmpAddr, peers[id]->address) && ::connect ( peers[id]->socketfd, (sockaddr*)&tmpAddr , sizeof(tmpAddr)) != -1)
                {
                    peers[id]->connected = true;

                    this->getCallback()->onPeerConnect(id);
                }
            }

        if (!peers[id]->connected)
        {
#ifdef WINDOWS
            if (errno != WSAECONNREFUSED)
#else
            if (errno != ECONNREFUSED)
#endif
                perror("connect");

            callback->onConnectionFailed(id, ENI_ConnectionDenied);

            this->removePeer(id);
        }
    }

    void CTCPPort::processIncomingConnection()
    {
        s32 socketfd = -1;
        SNetworkAddress addr;

        if (this->isIPv6())
        {
            sockaddr_in6 tmpAddr;
            socklen_t tmpSize = sizeof(tmpAddr);

            socketfd = accept( listenfd, (sockaddr*)&tmpAddr, &tmpSize);

            if (socketfd == -1)
                perror("accept");
            else
                getNetworkAddress6(tmpAddr, addr);
        }
        else
        {
            sockaddr_in tmpAddr;
            socklen_t tmpSize = sizeof(tmpAddr);

            socketfd = accept( listenfd, (sockaddr*)&tmpAddr, &tmpSize);

            if (socketfd == -1)
                perror("accept");
            else
                getNetworkAddress(tmpAddr, addr);
        }

        if (socketfd >= 0)
        {
            if (this->getCallback()->onConnectRequest(addr))
            {
                bool done = false;

                for (u16 u = 0; (u < maxPeers) && !done; u++)
                {
                    if (!peers[u])
                    {
                        peers.setValue(u, new STCPPeer(addr, socketfd));
                        peers[u]->connected = true;
                        this->addPeer( PeerID(u) );

                        this->getCallback()->onPeerConnect( PeerID(u) );

                        //! Don't forget to quit the loop
                        done = true;
                    }
                }
            }
            else
                close(socketfd);
        }
    }

} //! end namespace onet

